{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Memory模块\n",
    "\n",
    "Memory模块是ERNIEbot Agent的一个组件，可以作为无状态的LLM的记忆能力的补充，从而提升对话系统中的多轮交互效果。Memory的机制可以理解为将过去的消息进行存储并传递给LLM，从而LLM可以复用和对用户的问题相关的信息进行回答。\n",
    "\n",
    "因此对于Memory模块的设计主要考虑两个维度：\n",
    "- 存储容量：Memory模块能存储的消息不是无限多的，其不能超过LLM的上下文窗口长度，因此当消息超过了窗口长度后，需要将不相干的消息进行删除。\n",
    "- 语义相关度：Memory模块的作用时用于回答用户问题，因此需要保证存储的消息尽量与问题相关。\n",
    "\n",
    "目前，从存储容量角度出发，ERNIE Bot Agent中建立了三个的Memory相关类：\n",
    "\n",
    "- `WholeMemory`：全量记忆功能，存储所有消息。\n",
    "- `SlidingWindowMemory`：滑窗截断消息，限制对话轮数。\n",
    "- `LimitTokensMemory`：Token限制memory，限制对话memory种的token数量。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用Memory模块\n",
    "\n",
    "下面我们展示怎么使用不同类型的memory模块：\n",
    "\n",
    "### 环境准备：\n",
    "参考[环境准备](https://github.com/PaddlePaddle/ERNIE-SDK/blob/develop/docs/modules/preparation.md)文档进行ERNIEbot Agent相关环境准备。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### WholeMemory\n",
    "WholeMemory是全量记忆，存储所有消息。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<HumanMessage role: 'user', content: '请将这个图片中的单词识别出来', token_count: 14>, <AIMessage role: 'assistant', content: '好，这个图片中的单词为meticulous', token_count: 11>, <HumanMessage role: 'user', content: '这个单词meticulous是什么意思呢？', token_count: 21>, <AIMessage role: 'assistant', content: '这个单词的意思是挑剔的，关注细节的', token_count: 16>]\n"
     ]
    }
   ],
   "source": [
    "from erniebot_agent.memory import AIMessage, HumanMessage, WholeMemory\n",
    "\n",
    "memory = WholeMemory()\n",
    "humanmessage = HumanMessage(\"请将这个图片中的单词识别出来\")\n",
    "aimessage = AIMessage(\"好，这个图片中的单词为meticulous\")  # Fake AIMessage\n",
    "memory.add_message(humanmessage)\n",
    "memory.add_message(aimessage)\n",
    "humanmessage = HumanMessage(\"这个单词meticulous是什么意思呢？\")\n",
    "aimessage = AIMessage(\"这个单词的意思是挑剔的，关注细节的\")  # Fake AIMessage\n",
    "memory.add_message(humanmessage)\n",
    "memory.add_message(aimessage)\n",
    "print(memory.get_messages())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上面的例子可以看到，我们实例化了WholeMemory，然后通过`memory.add_message(msg)`方法将消息添加到memory中，并通过`memory.get_messages()`方法获取所有消息。\n",
    "除了使用`memory.add_message(msg)`方法添加消息外，我们也可以通过`memory.add_messages(msgs)`方法批量添加消息。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<HumanMessage role: 'user', content: '请将这个图片中的单词识别出来', token_count: 14>, <AIMessage role: 'assistant', content: '好，这个图片中的单词为meticulous', token_count: 11>, <HumanMessage role: 'user', content: '这个单词meticulous是什么意思呢？', token_count: 21>, <AIMessage role: 'assistant', content: '这个单词的意思是挑剔的，关注细节的', token_count: 16>, <HumanMessage role: 'user', content: '请帮我把这个单词meticulous存储到单词本中', token_count: 25>, <AIMessage role: 'assistant', content: '好的，单词meticulous已经存储到单词本中', token_count: 14>, <HumanMessage role: 'user', content: '请问现在我的单词本中都有什么单词呢？', token_count: 18>, <AIMessage role: 'assistant', content: '单词中目前有单词：meticulous', token_count: 9>]\n"
     ]
    }
   ],
   "source": [
    "messages = [\n",
    "    HumanMessage(\"请帮我把这个单词meticulous存储到单词本中\"),\n",
    "    AIMessage(\"好的，单词meticulous已经存储到单词本中\"),\n",
    "    HumanMessage(\"请问现在我的单词本中都有什么单词呢？\"),\n",
    "    AIMessage(\"单词中目前有单词：meticulous\"),\n",
    "]\n",
    "\n",
    "memory.add_messages(messages)\n",
    "print(memory.get_messages())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们使用`memory.add_messages(messages)`添加了一组消息，然后使用`memory.get_messages()`来获取所有消息，可以发现所有的message都正确加入。\n",
    "\n",
    "如果在对话结束想要删除所有的消息记录，我们可以使用`memory.clear_chat_history()`方法。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    }
   ],
   "source": [
    "memory.clear_chat_history()\n",
    "print(memory.get_messages())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 其他带截断功能的memory\n",
    "ERNIEbot Agent中还提供了`SlidingWindowMemory`和`LimitTokensMemory`两个memory模块，它们分别对应于滑动窗口截断记忆和token数量限制截断记忆。这两类阶段记忆的原理是类似的，即在加入消息时根据不同限制条件对历史消息进行删除。\n",
    "\n",
    "#### SlidingWindowMemory\n",
    "`SlidingWindowMemory(max_round, retained_round)`是滑动窗口截断记忆，在memory中存储固定轮数的消息。其中`max_round`表示memory中消息最多存储的轮数，`retained_round`表示在memory中会保留的初始消息的轮数，用于初始消息比较重要的场景。\n",
    "\n",
    "#### LimitTokensMemory\n",
    "`LimitTokensMemory(max_token_limit)`是token数量限制截断记忆，在memory中存储固定数量的token。其中`max_token_limit`表示memory中最多存储的消息的token数量，超过这个限制后，从头开始对消息进行删除。\n",
    "\n",
    "下面我们给出两段代码示例，分别展示`SlidingWindowMemory`和`LimitTokensMemory`的使用效果。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    HumanMessage(\"请帮我把这个单词meticulous存储到单词本中\"),\n",
    "    AIMessage(\"好的，单词meticulous已经存储到单词本中\"),\n",
    "    HumanMessage(\"请问现在我的单词本中都有什么单词呢？\"),\n",
    "    AIMessage(\"单词中目前有单词：meticulous\"),\n",
    "    HumanMessage(\"我想对单词本中的单词全部打印出对应的中文含义用于记忆\"),\n",
    "    AIMessage(\"好的，单词本中包括：meticulous的意思是挑剔的，关注细节的\"),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "裁剪前的消息为： [<HumanMessage role: 'user', content: '请帮我把这个单词meticulous存储到单词本中', token_count: 25>, <AIMessage role: 'assistant', content: '好的，单词meticulous已经存储到单词本中', token_count: 14>, <HumanMessage role: 'user', content: '请问现在我的单词本中都有什么单词呢？', token_count: 18>, <AIMessage role: 'assistant', content: '单词中目前有单词：meticulous', token_count: 9>, <HumanMessage role: 'user', content: '我想对单词本中的单词全部打印出对应的中文含义用于记忆', token_count: 26>, <AIMessage role: 'assistant', content: '好的，单词本中包括：meticulous的意思是挑剔的，关注细节的', token_count: 21>]\n",
      "裁剪后的消息为： [<HumanMessage role: 'user', content: '请问现在我的单词本中都有什么单词呢？', token_count: 18>, <AIMessage role: 'assistant', content: '单词中目前有单词：meticulous', token_count: 9>, <HumanMessage role: 'user', content: '我想对单词本中的单词全部打印出对应的中文含义用于记忆', token_count: 26>, <AIMessage role: 'assistant', content: '好的，单词本中包括：meticulous的意思是挑剔的，关注细节的', token_count: 21>]\n"
     ]
    }
   ],
   "source": [
    "from erniebot_agent.memory import SlidingWindowMemory\n",
    "\n",
    "memory = SlidingWindowMemory(max_round=2, retained_round=0)\n",
    "memory.add_messages(messages)\n",
    "print(\"裁剪前的消息为：\", messages)\n",
    "print(\"裁剪后的消息为：\", memory.get_messages())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The token count of the message has been set before\n",
      "The token count of the message has been set before\n",
      "The token count of the message has been set before\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "裁剪前的消息为： [<HumanMessage role: 'user', content: '请帮我把这个单词meticulous存储到单词本中', token_count: 25>, <AIMessage role: 'assistant', content: '好的，单词meticulous已经存储到单词本中', token_count: 14>, <HumanMessage role: 'user', content: '请问现在我的单词本中都有什么单词呢？', token_count: 18>, <AIMessage role: 'assistant', content: '单词中目前有单词：meticulous', token_count: 9>, <HumanMessage role: 'user', content: '我想对单词本中的单词全部打印出对应的中文含义用于记忆', token_count: 26>, <AIMessage role: 'assistant', content: '好的，单词本中包括：meticulous的意思是挑剔的，关注细节的', token_count: 21>]\n",
      "裁剪后的消息为： [<HumanMessage role: 'user', content: '请帮我把这个单词meticulous存储到单词本中', token_count: 25>, <AIMessage role: 'assistant', content: '好的，单词meticulous已经存储到单词本中', token_count: 14>, <HumanMessage role: 'user', content: '我想对单词本中的单词全部打印出对应的中文含义用于记忆', token_count: 26>, <AIMessage role: 'assistant', content: '好的，单词本中包括：meticulous的意思是挑剔的，关注细节的', token_count: 21>]\n"
     ]
    }
   ],
   "source": [
    "from erniebot_agent.memory import SlidingWindowMemory\n",
    "\n",
    "memory = SlidingWindowMemory(max_round=2, retained_round=1)\n",
    "memory.add_messages(messages)\n",
    "print(\"裁剪前的消息为：\", messages)\n",
    "print(\"裁剪后的消息为：\", memory.get_messages())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上述两个例子我们分别展示SlidingWindowMemory保留一轮消息，和同时保留首轮消息的效果。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The token count of the message has been set before\n",
      "The token count of the message has been set before\n",
      "The token count of the message has been set before\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "裁剪前的消息为： [<HumanMessage role: 'user', content: '请帮我把这个单词meticulous存储到单词本中', token_count: 25>, <AIMessage role: 'assistant', content: '好的，单词meticulous已经存储到单词本中', token_count: 14>, <HumanMessage role: 'user', content: '请问现在我的单词本中都有什么单词呢？', token_count: 18>, <AIMessage role: 'assistant', content: '单词中目前有单词：meticulous', token_count: 9>, <HumanMessage role: 'user', content: '我想对单词本中的单词全部打印出对应的中文含义用于记忆', token_count: 26>, <AIMessage role: 'assistant', content: '好的，单词本中包括：meticulous的意思是挑剔的，关注细节的', token_count: 21>]\n",
      "裁剪后的消息为： [<AIMessage role: 'assistant', content: '单词中目前有单词：meticulous', token_count: 9>, <HumanMessage role: 'user', content: '我想对单词本中的单词全部打印出对应的中文含义用于记忆', token_count: 26>, <AIMessage role: 'assistant', content: '好的，单词本中包括：meticulous的意思是挑剔的，关注细节的', token_count: 21>]\n"
     ]
    }
   ],
   "source": [
    "from erniebot_agent.memory import LimitTokensMemory\n",
    "\n",
    "memory = LimitTokensMemory(max_token_limit=60)\n",
    "memory.add_messages(messages)\n",
    "print(\"裁剪前的消息为：\", messages)\n",
    "print(\"裁剪后的消息为：\", memory.get_messages())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "细心的读者可以观察到，我们在上面的裁剪中，似乎不仅需要满足token限制的条件，因为输出的消息长度本来为两条也可以满足token数量限制要求。这是因为我们还有一个LLM输入需要为奇数条消息的限制条件，由于三条消息超过了token限制，因此最后只剩下了一条消息。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 总结\n",
    "综上，我们展示了Memory的设计出发点，并展示了三种Memory的使用方法和裁剪效果。\n",
    "\n",
    "可以看到，我们讲到了两种memory的设计出发点，但是并未考虑第二点语义信息，因此后续我们还会持续补充更多的memory满足不同应用和效果的需求。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
